#include <websocketpp/config/asio_no_tls.hpp>
#include <websocketpp/server.hpp>

#include <unordered_set>
#include <cstring>
#include <fileapi.h>
#include "head.h"
#include "DealShm.h"
#include "base64.h"
// 进程快照相关
#include <TlHelp32.h>	
// 显示文件信息
#include <sys/stat.h>	
#pragma comment(lib, "version.lib")


#define MESSAGEBOXA 1
#define MESSAGEBOXW 2
#define CREATEFILEA 3
#define WRITEFILE 4
#define READFILE 5
#define OPENFILE 6
#define HEAPALLOC 7
#define HEAPCREATE 8
#define HEAPDESTORY 9
#define HEAPFREE 10
#define REGCREATEKEYEXA 11
#define REGSETVALUEEXA 12
#define REGCLOSEKEY 13
#define REGOPENKEYEXA 14
#define REGDELETEVALUEA 15
#define SOCKET 16
#define BIND 17
#define SEND 18
#define CONNECT 19
#define MEMCPY 20
#define RECV 21



using namespace std;
typedef websocketpp::server<websocketpp::config::asio> server;

using websocketpp::lib::placeholders::_1;
using websocketpp::lib::placeholders::_2;
using websocketpp::lib::bind;

// pull out the type of messages sent by our config
typedef server::message_ptr message_ptr;

// 共享内存相关
HANDLE hSemap[2];                                           // 信号量句柄
WCHAR nameSemap[2][10] = { L"dbgSemap", L"dbgHeap" };       // 信号量名称
HANDLE hMap[2];                                             // 内存映射对象句柄
WCHAR nameShm[2][10] = { L"dbgShm", L"dgbHeap" };           // 共享内存名称
LPVOID pShm[2];                                             // 共享内存指针
func_info* shm;                                             // 函数信息结构体指针
heap_func_info* shm_heap;                                   // 堆函数信息结构体指针
std::mutex mtx[2];                                          // 共享内存读写锁

char outputStr[1000];                       // 专门暂存调试信息的
WCHAR outputWStr[1000] = { 0 };             // 专门暂存调试信息的

// 进程信息相关
DWORD processId;
char processJson[1000];


// 异常处理相关
unordered_set<string> folderSet;

int is_hooking = 0;
int is_sending = 0;
template<class T>
int JsonParserAndDecode(T inf, char* buffer, int bufferSize, uint8* decode, int codeSize)
{
    int type = 0;
    decltype(inf) info=(decltype(inf))inf;
    sscanf_s(info->argName[0], "%d", &type);
    char argJson[1000]=R"({)", st[1000];
    for (int i = 1; i <= info->argNum; i++)
    {
        char signleArg[200];
        sprintf_s(signleArg, R"("%s":"%s")", info->argName[i], info->argValue[i]);
        strcat_s(argJson, sizeof(argJson), signleArg);
        if (i != info->argNum)
        {
            strcat_s(argJson, sizeof(argJson), ",");
        }
        else
        {
            strcat_s(argJson, sizeof(argJson), "}");
        }
        
    }
    sprintf_s(st, sizeof(st), "%d-%d-%d %02d:%02d %02ds%03dms",
        info->st.wYear, info->st.wMonth, info->st.wDay, info->st.wHour,
        info->st.wMinute, info->st.wSecond, info->st.wMilliseconds);
    sprintf_s(buffer, bufferSize, R"({"type":%d,"name":"%s","arg":%s,"st":"%s"})", 
        type, info->argValue[0], argJson, st);
    if (base64_encode(buffer, decode, codeSize) == -1)
    {
        strcpy_s(buffer, sizeof(buffer), "Json is too long");
        base64_encode(buffer, decode, codeSize);
    }
    return type;
}


typedef DWORD(WINAPI* MyGetMappedFileName)(HANDLE, LPVOID, LPTSTR, DWORD);
int GetFileNameByHandle(HANDLE hFile, LPSTR buff, DWORD bufSize)
{
    HANDLE hfilemap = CreateFileMapping(hFile, NULL, PAGE_READWRITE, NULL, NULL, NULL);      // 获取文件映射句柄
    if (INVALID_HANDLE_VALUE == hfilemap)
    {
        printf("CreateFileMapping error: %d\n", GetLastError());
        return 0;
    }

    LPVOID lpmap = MapViewOfFile(hfilemap, FILE_MAP_READ | FILE_MAP_WRITE, NULL, NULL, 0);    // 获取文件映射对象
    if (NULL == lpmap)
    {
        printf("MapViewOfFile error: %d\n", GetLastError());
        return 0;
    }

    // GetMappedFileName未定义
    MyGetMappedFileName GetMappedFileName = (MyGetMappedFileName)GetProcAddress(LoadLibraryA("psapi.dll"), "GetMappedFileNameA");
    if (GetMappedFileName == NULL)
    {
        printf("Get 'GetMappedFileName' FuncAddress error.\n");
        return 0;
    }
    char DosStr[MAX_PATH] = { 0 };
    DWORD length = GetMappedFileNameA(GetCurrentProcess(), lpmap, DosStr, bufSize);         // 获取Dos设备名
    if (0 == length)
    {
        printf("GetMappedFileName error: %d\n", GetLastError());
        return 0;
    }

    char DosPath[MAX_PATH] = { 0 };
    char DriverString[MAX_PATH] = { 0 };
    if (0 == GetLogicalDriveStringsA(MAX_PATH, DriverString))                            // 获取电脑的盘符列表
    {
        printf("GetLogicalDriveStrings error: %d", GetLastError());
        return 0;
    }

    char* p = (char*)DriverString;                                                    //p用来指向盘符
    do
    {
        *(p + 2) = '\0';                            //QuerDosDevice第一个参数必须是c:这种类型的，不能有\，所以把\抹掉
        if (!QueryDosDeviceA((LPCSTR)p, DosPath, MAX_PATH))                             // 将盘符映射成DOS设备名
        {
            printf("QueryDosDrive error: %d\n", GetLastError());
            return FALSE;
        }

        char* q = strstr(DosStr, DosPath);      //检测buff中是否有DosDevice中的DosPath，有的话，p指向的那个字串就是要的盘符
        if (q != 0)
        {
            sprintf_s(buff,bufSize, "%s%s", p, DosStr + strlen(DosPath));                           //找到之后应该把DosPath替换成盘符
            return 1;
        }
        p = p + 4;                                                                          //指针移动到DriverString的下一个盘符处
    } while (*p != 0);
    return 0;
}


void send_message(server* s, websocketpp::connection_hdl hdl, char* str)
{
    try {
        s->send(hdl, str, websocketpp::frame::opcode::text);
    }
    catch (websocketpp::exception const& e) {
        if (!strcmp(e.what(), "Bad Connection")) {
            is_hooking = 0;
            is_sending = 0;
            MessageBoxW(NULL,L"意外断连！请用Web控制台断开连接。",L"提示",NULL);
            char order[30];
            if (processId != 0) {
                sprintf_s(order, sizeof(order), "taskkill /PID %lu /F", processId);
                system(order);
            }
        }
        else {
            std::cout << "Echo failed because: "
                << "(" << e.what() << ")" << std::endl;
        }
    }
}


typedef struct  {
    server* s;
    websocketpp::connection_hdl hdl;
}action123;

template<class T>
int DealFunc(server* s, websocketpp::connection_hdl hdl, T inf)
{
    decltype(inf) info = (decltype(inf))inf;
    unsigned temp;
    HANDLE p;
    char buffer[1000];
    uint8 decode[2000];
    stringstream ss;
    int type = JsonParserAndDecode(info, buffer, sizeof(buffer), decode, sizeof(decode));
    send_message(s, hdl, (char*)decode);
    printf_s("%s", buffer);
    switch (type)
    {
    case WRITEFILE: {
        sscanf_s(info->argValue[1], "%p", &p);
        /*
        temp = strtoul(info->argValue[2], NULL, 16);
        GetFinalPathNameByHandle(p, outputWStr, sizeof(buffer), FILE_NAME_OPENED);
        cout << "文件名是：" << outputWStr <<endl;
        GetFinalPathNameByHandle(p, outputWStr, sizeof(buffer), VOLUME_NAME_DOS);
        cout << "路径和驱动号是：" << outputWStr << endl;*/
        break;
    }
    case CREATEFILEA: {
        //MessageBoxA(NULL,"test","test",NULL);
        temp = strtoul(info->argValue[2], NULL, 16);    // dwDesiredAccess
        HANDLE p = CreateFileA("b.txt", GENERIC_READ | GENERIC_WRITE | GENERIC_EXECUTE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, NULL, NULL);
        
        break;
    }
    // 堆异常处理已写入UI中
    case HEAPCREATE: {
        printf_s("HeapCreate\n");
        break;
    }
    case HEAPDESTORY: {
        printf_s("HeapDestroy\n");
        break;
    }
    case HEAPFREE: {
        printf_s("HeapFree\n");
        break;
    }
    case REGOPENKEYEXA: {
        if (strstr(info->argValue[1], "SOFTWARE\Microsoft\Windows\CurrentVersion\Run")) {
            // TODO: 输出异常——程序试图修改注册表启动项
        }
        break;
    }
    default: {
        printf_s("Hooked\n");
        break;
    }
    }
    return 0;
}






int GetProcessPriority(HANDLE hProcess)
{
    switch (GetPriorityClass(hProcess))
    {
    case NORMAL_PRIORITY_CLASS:return 0;
    case  IDLE_PRIORITY_CLASS:return 1;
    case REALTIME_PRIORITY_CLASS:return 2;
    case HIGH_PRIORITY_CLASS:return 3;
    case ABOVE_NORMAL_PRIORITY_CLASS:return 4;
    case BELOW_NORMAL_PRIORITY_CLASS:return 5;
    default:return 6;
    }
}


void updateProcessInfo(server* s, websocketpp::connection_hdl hdl, PROCESS_INFORMATION pi, char* EXE) {
    // 输出处理的EXE的信息
    struct _stat buf;
    if (_stat(EXE, &buf))
        OutputError();
    processId = pi.dwProcessId;
    // 输出新创建的进程信息
    PROCESSENTRY32 PrE;
    PrE.dwSize = sizeof(PROCESSENTRY32); // 设置大小
    HANDLE pi2 = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);   // 对所有的进程快照
    if (pi2 != INVALID_HANDLE_VALUE) {
        int flag = Process32First(pi2, &PrE);
        while (flag)
        {
            if (PrE.th32ProcessID == pi.dwProcessId) {
                // 进程名称、ID、优先级
                char priority[7][20] = { "NORMAL", "IDLE" , "REALTIME", "HIGH", "ABOVENORMAL", "BELOWNORMAL", "NULL" };
                sprintf_s(processJson, sizeof(processJson), R"({"name": "%S", "id": %d,"priority": "%s", "thread":{"handle": "%p", "id": %d}, "EXE":{"size":%d,"drive":"%c"}})",
                    PrE.szExeFile, PrE.th32ProcessID, priority[GetProcessPriority(pi.hProcess)], pi.hThread, pi.dwThreadId, buf.st_size, buf.st_dev + 'A');
                //HMODULE hMod[100];
                //DWORD cbNeeded;
                //HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, PrE.th32ProcessID);
                //if (EnumProcessModules(hProcess, hMod, sizeof(hMod), &cbNeeded))
                //{
                //    //moduleNum = cbNeeded / sizeof(HMODULE); 
                //    //DWORD dw = GetModuleFileNameEx(hProcess, hMod, GetBuffer(MAX_PATH), MAX_PATH);
                //}
                //else {
                //    //OutputError();
                //}
                printf_s("%s\n", processJson);
                uint8 decode[2000];
                if (base64_encode(processJson, decode, 2000) == -1)
                {
                    strcpy_s(processJson, sizeof(processJson), R"({"err":"进程信息Json太长了。"})");
                    base64_encode(processJson, decode, 2000);
                }
                send_message(s, hdl, (char*)decode);
                break;
            }
            flag = Process32Next(pi2, &PrE);
        }
    }
    else {
        printf_s("进程快照失败！\n");
        exit(-1);
    }
    //CloseHandle(pi2);


    return;
}

typedef struct {
    server* s;
    websocketpp::connection_hdl hdl;
    WCHAR DirPath[MAX_PATH + 1];
    char DLLPath[MAX_PATH + 1];
    WCHAR EXE[MAX_PATH + 1];
}path_param;
path_param path;
HANDLE hThreadHook;

DWORD WINAPI ThreadHook(LPVOID lpParam)
{
    path_param* path = (path_param*)lpParam;
    STARTUPINFO si;
    PROCESS_INFORMATION pi;
    ZeroMemory(&si, sizeof(STARTUPINFO));
    ZeroMemory(&pi, sizeof(PROCESS_INFORMATION));
    si.cb = sizeof(STARTUPINFO);
    WCHAR DirPath[MAX_PATH + 1];
    wcscpy_s(DirPath, MAX_PATH, path->DirPath);                   // Dll的文件夹
    char DLLPath[MAX_PATH + 1];    // Dll的地址
    strcpy_s(DLLPath, MAX_PATH, path->DLLPath);
    WCHAR EXE[MAX_PATH + 1] = { 0 };
    wcscpy_s(EXE, MAX_PATH, path->EXE);          // 需要注入程序的地址

    if (DetourCreateProcessWithDllEx(EXE, NULL, NULL, NULL, TRUE,
        CREATE_DEFAULT_ERROR_MODE | CREATE_SUSPENDED, NULL, DirPath, &si, &pi,
        DLLPath, NULL))                             // Question: 为什么这里用Ex？(就是扩展的函数而已)
    {
        char exe_path[MAX_PATH+1];
        sprintf_s(exe_path, MAX_PATH, "%ws", EXE);
        updateProcessInfo(path->s, path->hdl, pi, exe_path);
        
        if (is_hooking) {
            ResumeThread(pi.hThread);
        }
    }
    else
    {
        is_hooking = 0;
        is_sending = 0;
        char error[100];
        uint8 decode[200];
        strcpy_s(processJson, sizeof(processJson), R"({"err":"进程打开失败!"})");
        base64_encode(processJson, decode, 200);
        send_message(path->s, path->hdl, (char*)decode);
        sprintf_s(error, "%d", GetLastError());     //MessageBoxA(NULL, error, NULL, NULL);
    }
    return 0;
}

//int nowType = 0;
//char lastType[3] = "";
//SYSTEMTIME lastSt;


DWORD WINAPI ThreadSendMessage(LPVOID lpParam)
{
    action123* p = (action123*)lpParam;
    server* s = p->s;
    websocketpp::connection_hdl hdl = p->hdl;
    while (is_hooking)
    {
        // TODO: 无法避免短时间内重复填入覆盖的问题；同时，会导致信号量紊乱、过多，重复输出最后填入的内容
        if (WaitForSingleObject(hSemap[0], 10) == WAIT_OBJECT_0  && is_sending) {
            func_info info;
            // 共享内存管理
            memcpy_s(&info, sizeof(info), shm, sizeof(info));
            // 处理不同的hook信息
            //DWORD dwThreadId = 0; // 一点失败的尝试
            //HANDLE hThreadDealFunc = ThreadCreate((char*)"DealFunc", ThreadDealFunc, dwThreadId, &info);
            DealFunc(s, hdl, &info);
        }
    }
    return 0;
}

DWORD WINAPI ThreadHeapMessage(LPVOID lpParam)
{
    action123* p = (action123*)lpParam;
    server* s = p->s;
    websocketpp::connection_hdl hdl = p->hdl;
    while (is_hooking)
    {
        // TODO: 无法避免短时间内重复填入覆盖的问题；同时，会导致信号量紊乱、过多，重复输出最后填入的内容
        if (WaitForSingleObject(hSemap[1], 10) == WAIT_OBJECT_0 && is_sending) {
            heap_func_info info;
            // 共享内存管理
            memcpy_s(&info, sizeof(info), shm_heap, sizeof(info));
            // 处理不同的hook信息
            //DWORD dwThreadId = 0; // 一点失败的尝试
            //HANDLE hThreadDealFunc = ThreadCreate((char*)"DealFunc", ThreadDealFunc, dwThreadId, &info);
            DealFunc(s, hdl, &info);
        }
    }
    return 0;
}


int JsonGetValue(char* json, char* name, char* valueBuffer, int bufferSize)
{
    std::string str = json;
    int i = str.find(name);
    if (i<str.length())
    {
        for (; str[i] != ':'; i++);
        for (; str[i] != '\"'; i++);
        i++;
        int j = i, num = 0;
        for (; str[j] != '\"'; j++)
        {
            str[num++] = str[j];
            if (str[j] == '\\')j++;
        }
        str = str.substr(0, num);
    }
    else
    {
        printf_s("Something Wrong.\n");
        return -1;
    }
    strcpy_s(valueBuffer, bufferSize, _strdup(str.c_str()));
    
    return 0;
}



// Define a callback to handle incoming messages
void on_message(server* s, websocketpp::connection_hdl hdl, message_ptr msg) {

    setlocale(LC_ALL, "chs");       // 保证中文正常输出
    string strMsg = msg->get_payload(); // hdl是什么

    char input[1000], buffer2[1024];
    uint8 buffer[1024];
    base64_decode((const uint8*)strMsg.c_str(), strMsg.length(), buffer2);

    printf_s("%s\n", buffer2);

    // check for a special command to instruct the server to stop listening so
    // it can be cleanly exited.
    if (!strcmp(buffer2, "stop hook")) {
        is_hooking = 0;
        is_sending = 0;
        char order[30];
        if (processId)
        {
            sprintf_s(order, sizeof(order), "taskkill /PID %lu /F", processId);
            system(order);
        }
        return;
    }
    if (!strcmp(buffer2, "pause")) {
        is_sending = 0;
        return;
    }

    if (!strcmp(buffer2, "start sending")) {
        is_sending = 1;
        return;
    }

    if (buffer2[0] == '{')
    {
        char exe_path[MAX_PATH] = { 0 }, now_path[MAX_PATH] = {0};
        path.s = s;
        path.hdl = hdl;
        // 获取被hook的程序的绝对路径
        JsonGetValue(buffer2, (char*)"exePath", exe_path, MAX_PATH);
        mbstowcs_s(NULL, path.EXE, MAX_PATH, exe_path, _TRUNCATE);
        //wprintf_s(L"%s\n", path.EXE);
        // 获取服务端当前的绝对路径
        JsonGetValue(buffer2, (char*)"hookPath", now_path, MAX_PATH);
        int cut = 0;
        char is_fan = '\\';
        for (int i = 0; now_path[i]; i++)if (now_path[i] == '\\' || now_path[i] == '/') { cut = i; is_fan = now_path[i]; }
        now_path[cut] = 0;
        mbstowcs_s(NULL, path.DirPath, MAX_PATH, now_path, _TRUNCATE);
        // 获取注射的DLL的绝对路径
        sprintf_s(path.DLLPath, MAX_PATH, "%s%cinnjectDll.dll", now_path, is_fan);       
    }
    if (!strcmp(buffer2, "start hook") && !is_hooking) {
        
        HANDLE hThreadHook = ThreadCreate((char*)"Hook", ThreadHook, &path);
        
        action123 p;
        p.s = s;
        p.hdl = hdl;
        is_hooking = 1;
        is_sending = 1;
        if (is_sending)
        {
            HANDLE hThreadDealFunc = ThreadCreate((char*)"SendMessage", ThreadSendMessage, &p);
            HANDLE hThreadDealHeap = ThreadCreate((char*)"HeapMessage", ThreadHeapMessage, &p);
            CloseHandle(hThreadDealFunc);
            CloseHandle(hThreadDealHeap);
        }
        CloseHandle(hThreadHook);
        return;
    }

    if (!strcmp(buffer2, "close server"))
    {
        s->stop_listening();
        return;
    }
    
    

}

DWORD WINAPI ThreadServer(LPVOID lpParam) {
    // 初始化共享内存相关参数
    shm = (func_info*)InitShm(hSemap[0], hMap[0], pShm[0], nameSemap[0], nameShm[0]);
    shm_heap = (heap_func_info*)InitShm(hSemap[1], hMap[1], pShm[1], nameSemap[1], nameShm[1]);

    // Create a server endpoint
    server echo_server;

    try {
        // Set logging settings
        echo_server.set_access_channels(websocketpp::log::alevel::all);
        echo_server.clear_access_channels(websocketpp::log::alevel::frame_payload);

        // Initialize Asio
        echo_server.init_asio();

        // Register our message handler
        echo_server.set_message_handler(bind(&on_message, &echo_server, ::_1, ::_2));

        // Listen on port 8092
        echo_server.listen(8092);

        // Start the server accept loop
        echo_server.start_accept();

        // Start the ASIO io_service run loop
        echo_server.run();
    }
    catch (websocketpp::exception const& e) {
        std::cout << e.what() << std::endl;
    }
    catch (...) {
        std::cout << "other exception" << std::endl;
    }

    DestoryShm(hMap[0], pShm[0]);
    DestoryShm(hMap[1], pShm[1]);
}